package main

import (
	"context"
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	_ "github.com/go-sql-driver/mysql"
	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
)

// HTTPConfig HTTP服务器配置
type HTTPConfig struct {
	EnableCORS        bool
	CORSOrigins       []string
	EnableHealthCheck bool
	HealthCheckPath   string
	RequestTimeout    time.Duration
	EnableLogging     bool
}

// DatabaseConfig 数据库配置结构
type DatabaseConfig struct {
	Host            string
	Port            string
	User            string
	Password        string
	Database        string
	MaxOpenConns    int
	MaxIdleConns    int
	ConnMaxLifetime time.Duration
	ConnMaxIdleTime time.Duration
}

// MySQLMCPServer MySQL MCP 服务器
type MySQLMCPServer struct {
	db     *sql.DB
	config DatabaseConfig
}

// QueryArgs 查询参数结构
type QueryArgs struct {
	Query string `json:"query" jsonschema:"required,description=SQL查询语句"`
}

// ListTablesArgs 列表表参数
type ListTablesArgs struct {
	Pattern string `json:"pattern" jsonschema:"description=表名匹配模式（可选）"`
}

// DescribeTableArgs 描述表参数
type DescribeTableArgs struct {
	TableName string `json:"table_name" jsonschema:"required,description=要描述的表名"`
}

// CreateTableArgs 创建表参数
type CreateTableArgs struct {
	Query string `json:"query" jsonschema:"required,description=CREATE TABLE SQL语句"`
}

// InsertDataArgs 插入数据参数
type InsertDataArgs struct {
	Query string `json:"query" jsonschema:"required,description=INSERT SQL语句"`
}

// UpdateDataArgs 更新数据参数
type UpdateDataArgs struct {
	Query string `json:"query" jsonschema:"required,description=UPDATE SQL语句"`
}

// DeleteDataArgs 删除数据参数
type DeleteDataArgs struct {
	Query string `json:"query" jsonschema:"required,description=DELETE SQL语句"`
}

// ShowDatabasesArgs 显示数据库参数
type ShowDatabasesArgs struct{}

// ShowIndexesArgs 显示索引参数
type ShowIndexesArgs struct {
	TableName string `json:"table_name" jsonschema:"required,description=要显示索引的表名"`
}

// ServerConfig 服务器配置结构
type ServerConfig struct {
	Mode string // "stdio", "http", "sse"
	Host string
	Port string
}

// main 主函数
func main() {
	// 从环境变量获取服务器配置
	serverConfig := ServerConfig{
		Mode: getEnvOrDefault("SERVER_MODE", "stdio"), // 默认使用stdio模式
		Host: getEnvOrDefault("SERVER_HOST", "localhost"),
		Port: getEnvOrDefault("SERVER_PORT", "8080"),
	}

	// 从环境变量获取HTTP配置
	httpConfig := HTTPConfig{
		EnableCORS:        getBoolEnv("HTTP_ENABLE_CORS", true),
		CORSOrigins:       getStringSliceEnv("HTTP_CORS_ORIGINS", []string{"*"}),
		EnableHealthCheck: getBoolEnv("HTTP_ENABLE_HEALTH_CHECK", true),
		HealthCheckPath:   getEnvOrDefault("HTTP_HEALTH_CHECK_PATH", "/health"),
		RequestTimeout:    getDurationEnv("HTTP_REQUEST_TIMEOUT", 30*time.Second),
		EnableLogging:     getBoolEnv("HTTP_ENABLE_LOGGING", true),
	}

	// 从环境变量获取数据库配置
	config := DatabaseConfig{
		Host:            getEnvOrDefault("MYSQL_HOST", "localhost"),
		Port:            getEnvOrDefault("MYSQL_PORT", "3306"),
		User:            getEnvOrDefault("MYSQL_USER", "root"),
		Password:        getEnvOrDefault("MYSQL_PASSWORD", "12345678"),
		Database:        getEnvOrDefault("MYSQL_DATABASE", "discuz_db"),
		MaxOpenConns:    getIntEnv("MYSQL_MAX_OPEN_CONNS", 25),
		MaxIdleConns:    getIntEnv("MYSQL_MAX_IDLE_CONNS", 25),
		ConnMaxLifetime: getDurationEnv("MYSQL_CONN_MAX_LIFETIME", 5*time.Minute),
		ConnMaxIdleTime: getDurationEnv("MYSQL_CONN_MAX_IDLE_TIME", 5*time.Minute),
	}

	// 验证必需的配置参数
	if config.User == "" {
		log.Fatal("错误：MYSQL_USER 环境变量是必需的")
	}
	if config.Database == "" {
		log.Fatal("错误：MYSQL_DATABASE 环境变量是必需的")
	}

	// 创建 MySQL MCP 服务器
	mysqlServer, err := NewMySQLMCPServer(config)
	if err != nil {
		log.Fatalf("创建 MySQL MCP 服务器失败: %v", err)
	}
	defer mysqlServer.Close()

	// 创建 MCP 服务器
	s := server.NewMCPServer(
		"MySQL MCP Server",
		"1.0.0",
		server.WithToolCapabilities(true),
		server.WithRecovery(),
	)

	// 注册工具
	mysqlServer.registerTools(s)

	// 根据配置选择启动方式
	switch strings.ToLower(serverConfig.Mode) {
	case "http":
		startHTTPServer(s, serverConfig, httpConfig, mysqlServer)
	case "sse":
		startSSEServer(s, serverConfig)
	case "stdio":
		startStdioServer(s)
	default:
		log.Fatalf("不支持的服务器模式: %s，支持的模式: stdio, http, sse", serverConfig.Mode)
	}
}

// startStdioServer 启动stdio服务器（原始方式）
func startStdioServer(s *server.MCPServer) {
	log.Println("启动 MySQL MCP 服务器 (stdio模式)...")
	if err := server.ServeStdio(s); err != nil {
		log.Fatalf("stdio服务器错误: %v", err)
	}
}

// startHTTPServer 启动HTTP服务器
func startHTTPServer(s *server.MCPServer, config ServerConfig, httpConfig HTTPConfig, mysqlServer *MySQLMCPServer) {
	addr := fmt.Sprintf("%s:%s", config.Host, config.Port)

	// 创建可流式HTTP服务器
	httpServer := server.NewStreamableHTTPServer(s,
		server.WithEndpointPath("/mcp"),
		server.WithHeartbeatInterval(30*time.Second),
	)

	// 创建自定义路由器
	mux := http.NewServeMux()

	// 添加健康检查端点
	if httpConfig.EnableHealthCheck {
		mux.HandleFunc(httpConfig.HealthCheckPath, func(w http.ResponseWriter, r *http.Request) {
			if r.Method != http.MethodGet {
				http.Error(w, "方法不允许", http.StatusMethodNotAllowed)
				return
			}

			// 检查数据库连接
			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer cancel()

			healthStatus := map[string]interface{}{
				"status":    "healthy",
				"timestamp": time.Now().Unix(),
				"service":   "mysql-mcp-server",
				"version":   "1.0.0",
			}

			// 简单的数据库健康检查
			if err := mysqlServer.db.PingContext(ctx); err != nil {
				healthStatus["status"] = "unhealthy"
				healthStatus["database_error"] = err.Error()
				w.WriteHeader(http.StatusServiceUnavailable)
			}

			w.Header().Set("Content-Type", "application/json")
			json.NewEncoder(w).Encode(healthStatus)
		})
		log.Printf("健康检查端点: http://%s%s", addr, httpConfig.HealthCheckPath)
	}

	// 应用中间件到整个服务器
	var handler http.Handler = httpServer

	// 添加健康检查路由的处理
	if httpConfig.EnableHealthCheck {
		originalHandler := handler
		handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			if r.URL.Path == httpConfig.HealthCheckPath {
				mux.ServeHTTP(w, r)
			} else {
				originalHandler.ServeHTTP(w, r)
			}
		})
	}

	// 添加CORS中间件
	if httpConfig.EnableCORS {
		handler = corsMiddleware(handler, httpConfig.CORSOrigins)
	}

	// 添加日志中间件
	if httpConfig.EnableLogging {
		handler = loggingMiddleware(handler)
	}

	// 添加超时中间件
	handler = timeoutMiddleware(handler, httpConfig.RequestTimeout)

	// 创建HTTP服务器
	srv := &http.Server{
		Addr:         addr,
		Handler:      handler,
		ReadTimeout:  30 * time.Second,
		WriteTimeout: 30 * time.Second,
		IdleTimeout:  60 * time.Second,
	}

	log.Printf("启动 MySQL MCP HTTP 服务器在 %s", addr)
	log.Printf("MCP端点: http://%s/mcp", addr)
	if httpConfig.EnableCORS {
		log.Printf("CORS已启用，允许源: %v", httpConfig.CORSOrigins)
	}

	// 设置优雅关闭
	setupGracefulShutdown(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()

		log.Println("正在关闭HTTP服务器...")
		if err := srv.Shutdown(ctx); err != nil {
			log.Printf("HTTP服务器关闭错误: %v", err)
		}
	})

	// 启动服务器
	if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Fatalf("HTTP服务器错误: %v", err)
	}
}

// corsMiddleware CORS中间件
func corsMiddleware(next http.Handler, origins []string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		origin := r.Header.Get("Origin")

		// 检查来源是否被允许
		allowed := false
		for _, allowedOrigin := range origins {
			if allowedOrigin == "*" || allowedOrigin == origin {
				allowed = true
				break
			}
		}

		if allowed {
			w.Header().Set("Access-Control-Allow-Origin", origin)
		}

		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization, X-Requested-With")
		w.Header().Set("Access-Control-Max-Age", "86400")

		// 处理预检请求
		if r.Method == "OPTIONS" {
			w.WriteHeader(http.StatusOK)
			return
		}

		next.ServeHTTP(w, r)
	})
}

// loggingMiddleware 日志中间件
func loggingMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()

		// 创建响应记录器来捕获状态码
		recorder := &statusRecorder{ResponseWriter: w, status: 200}

		next.ServeHTTP(recorder, r)

		duration := time.Since(start)
		log.Printf("%s %s %d %v %s",
			r.Method,
			r.URL.Path,
			recorder.status,
			duration,
			r.RemoteAddr,
		)
	})
}

// statusRecorder 状态记录器
type statusRecorder struct {
	http.ResponseWriter
	status int
}

func (r *statusRecorder) WriteHeader(status int) {
	r.status = status
	r.ResponseWriter.WriteHeader(status)
}

// timeoutMiddleware 超时中间件
func timeoutMiddleware(next http.Handler, timeout time.Duration) http.Handler {
	return http.TimeoutHandler(next, timeout, "请求超时")
}

// startSSEServer 启动SSE服务器
func startSSEServer(s *server.MCPServer, config ServerConfig) {
	addr := fmt.Sprintf("%s:%s", config.Host, config.Port)

	// 创建SSE服务器
	sseServer := server.NewSSEServer(s,
		server.WithSSEEndpoint("/sse"),
		server.WithMessageEndpoint("/message"),
		server.WithKeepAlive(true),
		server.WithKeepAliveInterval(30*time.Second),
	)

	log.Printf("启动 MySQL MCP SSE 服务器在 %s", addr)
	log.Printf("SSE端点: http://%s/sse", addr)
	log.Printf("消息端点: http://%s/message", addr)

	// 设置优雅关闭
	setupGracefulShutdown(func() {
		ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
		defer cancel()

		log.Println("正在关闭SSE服务器...")
		if err := sseServer.Shutdown(ctx); err != nil {
			log.Printf("SSE服务器关闭错误: %v", err)
		}
	})

	// 启动服务器
	if err := sseServer.Start(addr); err != nil {
		log.Fatalf("SSE服务器错误: %v", err)
	}
}

// setupGracefulShutdown 设置优雅关闭
func setupGracefulShutdown(shutdownFunc func()) {
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt, syscall.SIGTERM)

	go func() {
		<-c
		log.Println("收到关闭信号，正在优雅关闭服务器...")
		shutdownFunc()
		os.Exit(0)
	}()
}

// NewMySQLMCPServer 创建新的 MySQL MCP 服务器
func NewMySQLMCPServer(config DatabaseConfig) (*MySQLMCPServer, error) {
	// 构建数据库连接字符串
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%s)/%s?parseTime=true&loc=Local",
		config.User, config.Password, config.Host, config.Port, config.Database)

	// 连接数据库
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		return nil, fmt.Errorf("连接数据库失败: %w", err)
	}

	// 配置连接池
	db.SetMaxOpenConns(config.MaxOpenConns)
	db.SetMaxIdleConns(config.MaxIdleConns)
	db.SetConnMaxLifetime(config.ConnMaxLifetime)
	db.SetConnMaxIdleTime(config.ConnMaxIdleTime)

	// 测试数据库连接
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	if err := db.PingContext(ctx); err != nil {
		db.Close()
		return nil, fmt.Errorf("数据库连接测试失败: %w", err)
	}

	log.Printf("成功连接到 MySQL 数据库 %s:%s/%s", config.Host, config.Port, config.Database)
	log.Printf("连接池配置: 最大连接数=%d, 最大空闲连接数=%d, 连接最大生命周期=%v, 连接最大空闲时间=%v",
		config.MaxOpenConns, config.MaxIdleConns, config.ConnMaxLifetime, config.ConnMaxIdleTime)

	return &MySQLMCPServer{
		db:     db,
		config: config,
	}, nil
}

// Close 关闭数据库连接
func (m *MySQLMCPServer) Close() error {
	if m.db != nil {
		return m.db.Close()
	}
	return nil
}

// registerTools 注册所有工具
func (m *MySQLMCPServer) registerTools(s *server.MCPServer) {
	// 注册查询工具
	s.AddTool(
		mcp.NewTool("execute_query",
			mcp.WithDescription("执行任意 SQL 查询（SELECT、INSERT、UPDATE、DELETE 等）"),
			mcp.WithString("query",
				mcp.Required(),
				mcp.Description("要执行的 SQL 查询语句"),
			),
		),
		m.handleExecuteQuery,
	)

	// 注册列出表工具
	s.AddTool(
		mcp.NewTool("list_tables",
			mcp.WithDescription("列出当前数据库中的所有表"),
			mcp.WithString("pattern",
				mcp.Description("表名匹配模式（可选，支持 LIKE 语法）"),
			),
		),
		m.handleListTables,
	)

	// 注册描述表工具
	s.AddTool(
		mcp.NewTool("describe_table",
			mcp.WithDescription("获取表的结构信息（列、类型、约束等）"),
			mcp.WithString("table_name",
				mcp.Required(),
				mcp.Description("要描述的表名"),
			),
		),
		m.handleDescribeTable,
	)

	// 注册显示数据库工具
	s.AddTool(
		mcp.NewTool("show_databases",
			mcp.WithDescription("显示所有可用的数据库"),
		),
		m.handleShowDatabases,
	)

	// 注册显示索引工具
	s.AddTool(
		mcp.NewTool("show_indexes",
			mcp.WithDescription("显示指定表的所有索引"),
			mcp.WithString("table_name",
				mcp.Required(),
				mcp.Description("要显示索引的表名"),
			),
		),
		m.handleShowIndexes,
	)

	// 注册表状态工具
	s.AddTool(
		mcp.NewTool("table_status",
			mcp.WithDescription("获取表的状态信息（大小、行数、引擎等）"),
		),
		m.handleTableStatus,
	)

	// 注册备份表工具
	s.AddTool(
		mcp.NewTool("backup_table",
			mcp.WithDescription("创建表的备份副本"),
			mcp.WithString("table_name",
				mcp.Required(),
				mcp.Description("要备份的表名"),
			),
			mcp.WithString("backup_name",
				mcp.Description("备份表名（可选，默认为原表名_backup_timestamp）"),
			),
		),
		m.handleBackupTable,
	)
}

// handleExecuteQuery 处理执行查询
func (m *MySQLMCPServer) handleExecuteQuery(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	query, err := request.RequireString("query")
	if err != nil {
		return mcp.NewToolResultError("查询参数是必需的"), nil
	}

	// 清理查询语句
	query = strings.TrimSpace(query)
	if query == "" {
		return mcp.NewToolResultError("查询语句不能为空"), nil
	}

	// 判断查询类型
	queryType := strings.ToUpper(strings.Split(query, " ")[0])

	switch queryType {
	case "SELECT", "SHOW", "DESCRIBE", "DESC", "EXPLAIN":
		return m.executeSelectQuery(query)
	case "INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER":
		return m.executeModifyQuery(query)
	default:
		return mcp.NewToolResultError(fmt.Sprintf("不支持的查询类型: %s", queryType)), nil
	}
}

// executeSelectQuery 执行查询语句
func (m *MySQLMCPServer) executeSelectQuery(query string) (*mcp.CallToolResult, error) {
	rows, err := m.db.Query(query)
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("查询执行失败: %v", err)), nil
	}
	defer rows.Close()

	// 获取列信息
	columns, err := rows.Columns()
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("获取列信息失败: %v", err)), nil
	}

	// 存储结果
	var results []map[string]interface{}

	for rows.Next() {
		values := make([]interface{}, len(columns))
		valuePtrs := make([]interface{}, len(columns))

		for i := range values {
			valuePtrs[i] = &values[i]
		}

		if err := rows.Scan(valuePtrs...); err != nil {
			return mcp.NewToolResultError(fmt.Sprintf("扫描行数据失败: %v", err)), nil
		}

		row := make(map[string]interface{})
		for i, col := range columns {
			val := values[i]
			if b, ok := val.([]byte); ok {
				row[col] = string(b)
			} else {
				row[col] = val
			}
		}
		results = append(results, row)
	}

	if err := rows.Err(); err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("行遍历错误: %v", err)), nil
	}

	// 格式化结果
	resultJSON, err := json.MarshalIndent(results, "", "  ")
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("结果序列化失败: %v", err)), nil
	}

	summary := fmt.Sprintf("查询成功执行，返回 %d 行结果", len(results))

	return mcp.NewToolResultText(fmt.Sprintf("%s\n\n结果:\n%s", summary, string(resultJSON))), nil
}

// executeModifyQuery 执行修改语句
func (m *MySQLMCPServer) executeModifyQuery(query string) (*mcp.CallToolResult, error) {
	result, err := m.db.Exec(query)
	if err != nil {
		return mcp.NewToolResultError(fmt.Sprintf("查询执行失败: %v", err)), nil
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		rowsAffected = -1
	}

	lastInsertId, err := result.LastInsertId()
	if err != nil {
		lastInsertId = -1
	}

	var response strings.Builder
	response.WriteString("查询成功执行\n")
	response.WriteString(fmt.Sprintf("影响的行数: %d\n", rowsAffected))

	if lastInsertId > 0 {
		response.WriteString(fmt.Sprintf("最后插入的ID: %d\n", lastInsertId))
	}

	return mcp.NewToolResultText(response.String()), nil
}

// handleListTables 处理列出表
func (m *MySQLMCPServer) handleListTables(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	args := request.GetArguments()
	pattern := ""
	if p, ok := args["pattern"].(string); ok {
		pattern = p
	}

	var query string
	if pattern != "" {
		query = fmt.Sprintf("SHOW TABLES LIKE '%s'", pattern)
	} else {
		query = "SHOW TABLES"
	}

	return m.executeSelectQuery(query)
}

// handleDescribeTable 处理描述表
func (m *MySQLMCPServer) handleDescribeTable(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	tableName, err := request.RequireString("table_name")
	if err != nil {
		return mcp.NewToolResultError("表名参数是必需的"), nil
	}

	// 获取表结构
	query := fmt.Sprintf("DESCRIBE %s", tableName)

	return m.executeSelectQuery(query)
}

// handleShowDatabases 处理显示数据库
func (m *MySQLMCPServer) handleShowDatabases(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	return m.executeSelectQuery("SHOW DATABASES")
}

// handleShowIndexes 处理显示索引
func (m *MySQLMCPServer) handleShowIndexes(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	tableName, err := request.RequireString("table_name")
	if err != nil {
		return mcp.NewToolResultError("表名参数是必需的"), nil
	}

	query := fmt.Sprintf("SHOW INDEX FROM %s", tableName)

	return m.executeSelectQuery(query)
}

// handleTableStatus 处理表状态
func (m *MySQLMCPServer) handleTableStatus(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	return m.executeSelectQuery("SHOW TABLE STATUS")
}

// handleBackupTable 处理备份表
func (m *MySQLMCPServer) handleBackupTable(ctx context.Context, request mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	tableName, err := request.RequireString("table_name")
	if err != nil {
		return mcp.NewToolResultError("表名参数是必需的"), nil
	}

	args := request.GetArguments()
	backupName := ""
	if b, ok := args["backup_name"].(string); ok {
		backupName = b
	}
	if backupName == "" {
		backupName = fmt.Sprintf("%s_backup_%d", tableName,
			getCurrentTimestamp())
	}

	query := fmt.Sprintf("CREATE TABLE %s AS SELECT * FROM %s", backupName, tableName)

	return m.executeModifyQuery(query)
}

// getEnvOrDefault 获取环境变量或默认值
func getEnvOrDefault(key, defaultValue string) string {
	if value := os.Getenv(key); value != "" {
		return value
	}
	return defaultValue
}

// getCurrentTimestamp 获取当前时间戳
func getCurrentTimestamp() int64 {
	return time.Now().Unix()
}

// getBoolEnv 获取布尔环境变量
func getBoolEnv(key string, defaultValue bool) bool {
	if value := os.Getenv(key); value != "" {
		return value == "true"
	}
	return defaultValue
}

// getStringSliceEnv 获取字符串切片环境变量
func getStringSliceEnv(key string, defaultValue []string) []string {
	if value := os.Getenv(key); value != "" {
		return strings.Split(value, ",")
	}
	return defaultValue
}

// getIntEnv 获取整数环境变量
func getIntEnv(key string, defaultValue int) int {
	if value := os.Getenv(key); value != "" {
		intValue, err := strconv.Atoi(value)
		if err == nil {
			return intValue
		}
	}
	return defaultValue
}

// getDurationEnv 获取持续时间环境变量
func getDurationEnv(key string, defaultValue time.Duration) time.Duration {
	if value := os.Getenv(key); value != "" {
		duration, err := time.ParseDuration(value)
		if err == nil {
			return duration
		}
	}
	return defaultValue
}
